# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import logging
import os

import numpy

from coin.database.base.serializer import Serializer
from coin.database.base import h5py_io


class FileNotFound(Exception):
  pass


class NumpySerializer(Serializer):
  def __init__(self):
    Serializer.__init__(self)

  def serialize_impl(self, cache_path, data, compress):
    if data.dtype is float:
      data[data == ''] = numpy.nan
      data[data == 'None'] = numpy.nan
    if data.dtype.type is numpy.unicode_:
      try:
        data = data.astype(numpy.string_)
      except Exception:
        logging.warning("Unicode to string conversion failed for %s.", cache_path)

    if compress is True:
      # Npz is deprecated.
      h5py_io.h5py_save(cache_path + ".hdf5", data, "lzf")
    else:
      numpy.save(cache_path, data, allow_pickle=False)

  def deserialize_impl(self, cache_path):
    if os.path.exists(cache_path + ".npy"):
      return numpy.load(cache_path + ".npy", mmap_mode="r")
    elif os.path.exists(cache_path + ".npz"):
      return numpy.load(cache_path + ".npz", mmap_mode="r")["arr_0"][0]
    elif os.path.exists(cache_path + ".hdf5"):
      return h5py_io.h5py_load(cache_path + ".hdf5")
    else:
      raise FileNotFound(cache_path)
